import openai
import re
from collections import deque
from contextlib import redirect_stdout
import io


class CodeExecutor:
    '''代码执行模块（带持久化环境）'''
    def __init__(self):
        self.locals = {}
        self.blacklist = ['open(', 'os.', 'sys.', 'subprocess.']  # 基础安全过滤
    
    def _safety_check(self, code):
        for forbidden in self.blacklist:
            if forbidden in code:
                return False, f'安全检查失败：检测到禁用关键字 {forbidden}'
        return True, ''
    
    def execute(self, code):
        # 执行安全检查
        is_safe, msg = self._safety_check(code)
        if not is_safe:
            return msg
        
        try:
            # 准备执行环境
            f = io.StringIO()
            with redirect_stdout(f):
                exec(code, self.locals, self.locals)  # 执行代码
            output = f.getvalue()
            return output.strip() or '代码执行成功（无控制台输出）'
        except Exception as e:
            return f'执行错误：{str(e)}'
class LLM:
    '''大语言模型集成模块'''
    def __init__(self):
        api_key = '<用户 API Key>'
        base_url = 'https://dashscope.aliyuncs.com/compatible-mode/v1'
        self.model = 'qwen-turbo-latest'
        self.client = openai.OpenAI(
            api_key=api_key,
            base_url=base_url,
        )

    def generate_response(self, messages):
        try:
            response = self.client.chat.completions.create(
                model=self.model,
                messages=messages,
                temperature=0.3,  # 降低随机性保证代码准确性
                max_tokens=2000
            )
            return response.choices[0].message.content.strip()
        except Exception as e:
            return f'生成响应时出错：{str(e)}'


class MemoryManager:
    '''记忆管理模块'''
    def __init__(self, max_history=5):
        self.history = deque(maxlen=max_history)

    def add_history(self, role, content):
        self.history.append({'role': role, 'content': content})

    def get_history(self):
        return list(self.history)


class CodeAgent:
    '''代码执行智能体'''
    def __init__(self):
        self.llm = LLM()
        self.memory = MemoryManager()
        self.executor = CodeExecutor()
        self.system_prompt = {
            'role': 'system',
            'content': '''你是一个专业代码助手，请遵守以下规则：
1. 仅生成符合要求的Python代码
2. 代码必须用三个反引号包裹
3. 优先使用标准库
4. 保证代码简洁高效
5. 不要包含解释性文本'''
        }
        self.code_pattern = re.compile(r'```python\n(.*?)\n```', re.DOTALL)

    def _build_messages(self, history, prompt):
        messages = [self.system_prompt] + history
        messages.append({'role': 'user', 'content': prompt})
        return messages

    def _extract_code(self, response):
        match = self.code_pattern.search(response)
        return match.group(1).strip() if match else None

    def process_input(self, user_input):
        # 维护对话历史
        self.memory.add_history('user', user_input)
        
        # 准备对话上下文
        history = self.memory.get_history()
        messages = self._build_messages(history, user_input)
        
        # 获取LLM响应
        response = self.llm.generate_response(messages)
        
        # 处理代码执行
        execution_result = ''
        code = self._extract_code(response)
        print(f'[code]\n{code}')
        if code:
            execution_result = self.executor.execute(code)
            final_response = f'执行结果：\n{execution_result}'
        else:
            final_response = '未检测到有效代码，请确认是否满足以下要求：\n1. 使用python代码块\n2. 正确使用三个反引号包裹'
        
        # 更新对话历史
        self.memory.add_history('assistant', final_response)
        return final_response


def main():
    agent = CodeAgent()
    examples = [
        '写一个打印Hello World的程序',
        '创建一个从1加到100的求和函数并调用',
        '写一个计算斐波那契数列前n项的函数',
        '帮我生成一个三维坐标点的类，包含距离计算方法',
        '刚才的代码执行结果是什么？'
    ]
    
    for query in examples:
        print(f'[用户] {query}')
        response = agent.process_input(query)
        print(f'[助手] {response}\n{"-"*60}')


if __name__ == '__main__':
    main()